#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
调用tds4py模块生成频散图的示例脚本
"""

import numpy as np
import matplotlib.pyplot as plt
import os
import sys

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

try:
    from TDS import TDS
except ImportError:
    print("无法导入tds4py模块，请确保已正确编译")
    sys.exit(1)

def main():
    # 创建TDS实例
    tds = TDS()
    
    # 检查Fortran模块是否可用
    if not tds.use_fortran:
        print("Fortran模块不可用，无法进行计算")
        return
    
    print("开始计算频散谱...")
    print(f"频率范围: {tds.source['mino']} - {tds.source['maxo']} Hz")
    print(f"速度范围: {tds.source['minv']} - {tds.source['maxv']} km/s")
    print(f"网格: {tds.source['no']} x {tds.source['nv']}")
    
    # 计算频散谱
    dispersion_result = tds.calculate_dispersion_spectrum()
    
    # 获取参数用于绘图
    no = tds.source['no']
    nv = tds.source['nv']
    mino = tds.source['mino']
    maxo = tds.source['maxo']
    minv = tds.source['minv']
    maxv = tds.source['maxv']
    
    # 创建频率和速度数组
    frequencies = np.linspace(mino, maxo, no)
    velocities = np.linspace(minv, maxv, nv)
    
    # 计算幅度谱
    amplitude_spectrum = np.abs(dispersion_result)
    
    # 绘制频散图
    plt.figure(figsize=(12, 8))
    plt.imshow(amplitude_spectrum.T, 
               extent=[frequencies[0], frequencies[-1], velocities[0], velocities[-1]], 
               aspect='auto', 
               origin='lower', 
               cmap='jet')
    plt.colorbar(label='幅度')
    plt.xlabel('频率 (Hz)')
    plt.ylabel('相速度 (km/s)')
    plt.title('理论频散谱')
    plt.tight_layout()
    
    # 保存图像
    plt.savefig('dispersion_spectrum.png', dpi=300, bbox_inches='tight')
    print("频散图已保存为 dispersion_spectrum.png")
    
    # 保存数据为文本文件（类似原始程序的.ds文件格式）
    np.savetxt('dispersion_spectrum.ds', amplitude_spectrum)
    print("频散数据已保存为 dispersion_spectrum.ds")
    
    # 显示图像
    plt.show()
    
    print("计算完成!")

if __name__ == "__main__":
    main()